import jsonpath
import re
import logging

logger = logging.getLogger(__name__)


def extract(resp, attr_name, pattern, index=0):
    if attr_name == "json":
        resp.json = resp.json()  # 把方法改为属性

    data = getattr(resp, attr_name)
    if pattern.startswith("$"):
        # 使用jsonpatn
        # print("使用jsonpatn")
        value_list = jsonpath.jsonpath(dict(data), pattern)
        if value_list:
            value = value_list[index]
            logger.info(f"提取的数据{value}")
            return value
        else:
            logger.error(f"数据提取失败，检查规则是否正确，pattern：{pattern}，resp.attr_name：{dict(data)}")
    else:
        # 使用正则
        # print("使用正则")
        value_list = re.findall(pattern, data)
        if value_list:
            value = value_list[index]
            logger.info(f"提取的数据{value}")
            return value
        else:
            logger.error(f"数据提取失败，检查规则是否正确，字段：{pattern}，数据：{data}")
